package main

import (
	"bytes"
	"crypto/cipher"
	"crypto/rand"
	"fmt"
	"log"

	"github.com/Hyperledger-TWGC/tjfoc-gm/sm2"
	"github.com/Hyperledger-TWGC/tjfoc-gm/sm3"
	"github.com/Hyperledger-TWGC/tjfoc-gm/sm4"
)

func main() {
	SM2()

	SM3()

	SM4()
}

func SM2() {
	priv, err := sm2.GenerateKey(rand.Reader) // 生成密钥对
	if err != nil {
		log.Fatal(err)
	}
	msg := []byte("Tongji Fintech Research Institute")
	pub := &priv.PublicKey
	ciphertxt, err := sm2.Encrypt(pub, msg, rand.Reader)
	if err != nil {
		log.Fatal(err)
	}
	fmt.Println("SM2加密密文是：", string(msg))
	fmt.Printf("SM2加密结果是：%x\n", ciphertxt)
	
	plaintxt, err := priv.Decrypt(nil, ciphertxt, nil)
	if err != nil {
		log.Fatal(err)
	}
	if !bytes.Equal(msg, plaintxt) {
		log.Fatal("原文不匹配")
	}

	sign, err := priv.Sign(rand.Reader, msg, nil)
	if err != nil {
		log.Fatal(err)
	}

	isok := priv.PublicKey.Verify(msg, sign)
	fmt.Printf("SM2 Verified: %v\n", isok)
}

func SM3() {
	data := "test"
	h := sm3.New()
	h.Write([]byte(data))
	sum := h.Sum(nil)
	fmt.Printf("SM3 digest value is: %x\n", sum)
}

func SM4() {
	// 128比特密钥
	key := []byte("1234567890abcdef")
	// 128比特iv
	iv := make([]byte, sm4.BlockSize)
	data := []byte("Tongji Fintech Research Institute")
	fmt.Println("SM4加密密文是：", string(data))
	ciphertxt, err := sm4Encrypt(key, iv, data)
	if err != nil {
		log.Fatal(err)
	}
	fmt.Printf("SM4加密结果: %x\n", ciphertxt)

	res, err := sm4Decrypt(key, iv, ciphertxt)
	if err != nil {
		log.Fatal(err)
	}
	fmt.Printf("SM4解密结果: %x\n", res)
	fmt.Println("SM4解密密文是：", string(res))
}

func sm4Encrypt(key, iv, plainText []byte) ([]byte, error) {
	block, err := sm4.NewCipher(key)
	if err != nil {
		return nil, err
	}
	blockSize := block.BlockSize()
	origData := pkcs5Padding(plainText, blockSize)
	blockMode := cipher.NewCBCEncrypter(block, iv)
	cryted := make([]byte, len(origData))
	blockMode.CryptBlocks(cryted, origData)
	return cryted, nil
}

func sm4Decrypt(key, iv, cipherText []byte) ([]byte, error) {
	block, err := sm4.NewCipher(key)
	if err != nil {
		return nil, err
	}
	blockMode := cipher.NewCBCDecrypter(block, iv)
	origData := make([]byte, len(cipherText))
	blockMode.CryptBlocks(origData, cipherText)
	origData = pkcs5UnPadding(origData)
	return origData, nil
}

// pkcs5填充
func pkcs5Padding(src []byte, blockSize int) []byte {
	padding := blockSize - len(src)%blockSize
	padtext := bytes.Repeat([]byte{byte(padding)}, padding)
	return append(src, padtext...)
}

func pkcs5UnPadding(src []byte) []byte {
	length := len(src)
	if length == 0 {
		return nil
	}
	unpadding := int(src[length-1])
	return src[:(length - unpadding)]
}


